{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载数据集A&B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载A和B测试集\n",
    "import os\n",
    "folder_path = r'Dataset/Testset_track_B/'\n",
    "if os.path.exists(folder_path) and os.path.isdir(folder_path):\n",
    "    print(\"数据文件夹存在\")\n",
    "else:\n",
    "    !wget --header=\"Host: drive.usercontent.google.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Cookie: __Secure-ENID=12.SE=Yd0Bj-CLJ14fnd4qzdJHmwUs4B5zz46UaPC1cPJigNqqFV9PtM2CYyBpSbCkOyzUwzlEdZ1nZFf-igtGi7wSdJ_gqQSfQfh84r9egqFQAy9-GKayCRbdQKdera-2mkpuIT-c64CyR9vfNojM3hxZ9Dej-dGvtxlGjal9ttEHybw; __gsas=ID=ae0421b9a34b478c:T=1710758437:RT=1710758437:S=ALNI_MZP13R9ZOHbCzC0rgHSMrGXj6GCsg; HSID=A-4I-ZudDNUIB6EKH; SSID=A7v_1v9un6xAwVNku; APISID=ctK8IbLjeuDUmgys/AFnMSLWt9KddceDI6; SAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; __Secure-1PAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; __Secure-3PAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; SID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_kzuBV1TvOhAIC8VF1e9fpgACgYKATQSARQSFQHGX2Mi8LXUwWoIwNCEPU8Sy3mXUxoVAUF8yKqGXVfjTGz9gQal7nwGr4Pl0076; __Secure-1PSID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_PDa-DzVmbdGFPyxMQpk9_QACgYKAewSARQSFQHGX2MiAeee4fn0OWglWZfAygqkyBoVAUF8yKp-Sfmtnueimxc-0QbJRF9I0076; __Secure-3PSID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_g9IrMeU98APBo9Stp6wEnAACgYKAQASARQSFQHGX2MiFWtc9ucONXnpxBzlRdudEhoVAUF8yKoeZwCpJDnjfAFjGssHSUGm0076; NID=515=GQhY9nKKFCx3qFDjE0MA4ubjWNdef6xCIY_RfWOPWKEtyfBN3nAUl8WHI2VczjNQ4rVkj1XBAY8WNWHXyqSK10CfT4FxsFlPzrHIJpeTtm1nWRNBd9AAfBKJHz4XpESszntVUTE_59RklZuKo0vk1poReVi2da1PZKC3CTKH2Ll3gB5xuB9wf4bmq8ylVUuIROPJczr0XnCuUHV3qLdBvgy9_870b6UwOq1iOlIxFQFm01EZ4pqF4q1Ub3QRSWpEMLh4LSZFpJ5O255R5OV7krmEdDvH_sHoTEPZAg2PoEpwAyGK6Xp9qcLIlldgx5-5V86N8Wtb93uTlQuA_CFXb5_2eP3bgeX8txwlJ5SrldVjg9ctzYtBU2RwJKTSvdHfIG7lpOkg6XlkvDOcJpR3DihT_OlqnPn7drCAJpvVDv29hZn5XPMXaSrNdbG64OJ9urJEw5odEwsLYkkpC1vmlUcuoo52S5f6RQu0Z8kZiV8iRW6XIqHsSmQHunVaxk6xWCStUg; __Secure-1PSIDTS=sidts-CjEB3EgAEtTS0OazynCofIH4RCBstiRP5flEcvYW3z4Fg9oGd5QOESDOZt1wO2iqUYHjEAA; __Secure-3PSIDTS=sidts-CjEB3EgAEtTS0OazynCofIH4RCBstiRP5flEcvYW3z4Fg9oGd5QOESDOZt1wO2iqUYHjEAA; SIDCC=AKEyXzVI6aMX8lSDja86Yts3FBAtBzPCzVNgaX5BCz78NWsWzlT3yFWKUV7ZE46SFzE1GiBI-cHdTw; __Secure-1PSIDCC=AKEyXzUo4NQAwqqPMxP2eye-MFEbZmBIm_sZqRU1amttg0YoQkc8ZKSNXdHl5jNCMEbhrUHhS9-K; __Secure-3PSIDCC=AKEyXzWf2lIdmDLeZKpXSi9GytVQb6XudrYiNUBA5gW952YuLh8kL6T3IbBlu8zOTfGEcdUp5O1R\" --header=\"Connection: keep-alive\" \"https://drive.usercontent.google.com/download?id=1JwR0Q1ArTg6c47EF2ZuIBpQwCPgXKrO2&export=download&authuser=0&confirm=t&uuid=dc3aa13c-c3a9-458f-983a-8586798cb635&at=APZUnTX25XMxi-z-3wBcgR93IGsL%3A1719235792953\" -c -O \"Dataset.zip\"\n",
    "    !unzip Dataset.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 如果需要训练C则需要解除下列注释"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/38e9adf0fce84527aad3558cc3e82d0e9a251aac4c934297afae9b74d9b3d1e9?responseContentDisposition=attachment%3B%20filename%3Dtrain_track_B.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-04T03%3A21%3A02Z%2F-1%2F%2Facd359add161bace603a52c7a268467406cb3c1889a7114bbb687de8002b55f6\" -c -O 'train_track_B.zip'\n",
    "# !wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/7877c2fd300345599ed3365feda50425f4857caa71bf4af9bf047fb08e35aa97?responseContentDisposition=attachment%3B%20filename%3Dmesh_B_0603.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-04T11%3A48%3A26Z%2F-1%2F%2Fec6cce492ba6afff841ef197860065742a7ba8220def02c1cbc311333aa993b2\" -c -O 'mesh_B_0603.zip'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !unzip -d train_track_B train_track_B.zip\n",
    "# !unzip -d mesh_B_0603 mesh_B_0603.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 安装必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!pip3 install open3d\n",
    "!pip3 install timm\n",
    "!pip3 install einops\n",
    "!pip3 install unzip\n",
    "!pip3 install wget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载最优权重A&B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trackA\n",
    "#  https://drive.google.com/file/d/1-Bo2wR4lJ1ZiHDF4EbhhgbE2ky9Dcuw9/view?usp=sharing\n",
    "!wget  --no-check-certificate 'https://docs.google.com/uc?export=download&id=1-Bo2wR4lJ1ZiHDF4EbhhgbE2ky9Dcuw9' -O best_weightsA.pt\n",
    "# trackB\n",
    "# # https://drive.google.com/file/d/1zADAmyG0ua5KHB_P_FRRFfbk7u9PlZCJ/view?usp=drive_link\n",
    "# !wget  --no-check-certificate 'https://docs.google.com/uc?export=download&id=1zADAmyG0ua5KHB_P_FRRFfbk7u9PlZCJ' -O best_weightsB.pt\n",
    "\n",
    "# https://drive.google.com/file/d/1XCPqRi-4uFdWdZgeWKhEbe16K07zb5ZG/view?usp=drive_link # 更新后的\n",
    "!wget  --no-check-certificate 'https://docs.google.com/uc?export=download&id=1XCPqRi-4uFdWdZgeWKhEbe16K07zb5ZG' -O best_weightsB.pt\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# trackC 加载数据 分别是A和B的模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 读取A的测试数据和训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据加载器\n",
    "# 加载测试集\n",
    "# 代码模型\n",
    "import numpy as np\n",
    "import open3d\n",
    "import os\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "# 按照顺序读取标签\n",
    "def read_order(path):\n",
    "    path = os.path.join(path, 'watertight_meshes.txt')\n",
    "    order = []\n",
    "    with open(path,mode='r') as f:\n",
    "        lines = f.readlines()\n",
    "        for line in lines:\n",
    "            # 需要清除换行符\n",
    "            order.append(line.strip())\n",
    "    return order        \n",
    "\n",
    "# 加载特征和标签的路径\n",
    "def load_train(path=r'Dataset/Training_data'):\n",
    "    order = read_order(path=path)\n",
    "    feature_path = []\n",
    "    label_path = []\n",
    "    f_base_path = os.path.join(path, 'Feature')\n",
    "    l_base_path = os.path.join(path, 'Label')\n",
    "    for o in order:\n",
    "        f_path = f_base_path + f\"/mesh_{o}.ply\"\n",
    "        l_path = l_base_path + f\"/press_{o}.npy\"\n",
    "        feature_path.append(f_path)\n",
    "        label_path.append(l_path)\n",
    "    return feature_path, label_path\n",
    "\n",
    "\n",
    "\n",
    "class MeshPressDataset(Dataset):\n",
    "    def __init__(self, feature_dir, press_dir):\n",
    "        # 特征和标签的文件夹\n",
    "        self.feature_dir = feature_dir\n",
    "        self.press_dir = press_dir\n",
    "        # 判读长度是否一致\n",
    "        assert len(self.feature_dir) == len(self.press_dir)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.feature_dir)\n",
    "    \n",
    "    # 创造一个迭代器获取数据集\n",
    "    def __getitem__(self, idx):\n",
    "        pcd = open3d.io.read_point_cloud(self.feature_dir[idx])\n",
    "        # 转换为np数据\n",
    "        vertices = np.array(pcd.points)\n",
    "        # 在此处进行裁剪压力\n",
    "        # print(self.label_paths[idx])\n",
    "        press = np.load(self.press_dir[idx])\n",
    "        press = np.concatenate((press[:16], press[112:]), axis=0)\n",
    "        # 确保数据长度一致\n",
    "        assert vertices.shape[0] == press.shape[0]\n",
    "\n",
    "        return torch.tensor(vertices, dtype=torch.float32), torch.tensor(press, dtype=torch.float32)\n",
    "\n",
    "\n",
    "\n",
    "def read_test(path=r'Dataset/Testset_track_A/Inference/'):\n",
    "    \"\"\"需要传入测试集文件夹路径\"\"\"\n",
    "    t_base_path = path\n",
    "    paths = os.listdir(t_base_path)\n",
    "    test_path = []\n",
    "    for p in paths:\n",
    "        t_path = t_base_path + p\n",
    "        test_path.append(t_path)\n",
    "    return test_path\n",
    "\n",
    "class TestDataset(Dataset):\n",
    "    def __init__(self, test_dir):\n",
    "        self.test_dir = test_dir\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.test_dir)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        pcd = open3d.io.read_point_cloud(self.test_dir[idx])\n",
    "        # 转换为np数据\n",
    "        vertices = np.array(pcd.points)\n",
    "\n",
    "        return torch.tensor(vertices, dtype=torch.float32), self.test_dir[idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 读取B的训练数据和测试数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 用于加载B赛道的数据集\n",
    "\"\"\"\n",
    "1- 首先获取存储数据的路径\n",
    "2- 分别加载训练集和验证集、最后加载测试集（if need test）\n",
    "\"\"\"\n",
    "\n",
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "from pathlib import Path\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "class MeshPressDatasetB(Dataset):\n",
    "    def __init__(\n",
    "            self,\n",
    "            train_mesh_dir = None,\n",
    "            train_centroid_dir = None,\n",
    "            test_mesh_dir = None,\n",
    "            n_train: int = 500,\n",
    "            val_ratio=0.2,\n",
    "            mode='train',\n",
    "            valid_val_indices=None):\n",
    "        Dataset.__init__(self)\n",
    "        \n",
    "        self.mode = mode\n",
    "        if self.mode == 'test':\n",
    "            self.all_test_mesh_path = [self.get_mesh_path_test(test_mesh_dir, i) for i in range(1, 51)]\n",
    "            self.all_test_centroid_path = [self.get_centroid_data_path_test(test_mesh_dir, i) for i in range(1, 51)]\n",
    "            self.all_test_area_path = [self.get_area_data_path_test(test_mesh_dir, i) for i in range(1, 51)]\n",
    "        else:\n",
    "            valid_train_indices = self.get_all_mesh_indices(train_mesh_dir)\n",
    "            random.shuffle(valid_train_indices)\n",
    "            val_num = int(n_train * val_ratio)\n",
    "            if valid_val_indices is not None:\n",
    "                valid_val_indices = valid_val_indices\n",
    "            else:\n",
    "                valid_val_indices = valid_train_indices[:val_num]\n",
    "            valid_train_indices = valid_train_indices[val_num:]\n",
    "            self.valid_val_indices = valid_val_indices\n",
    "            self.all_train_mesh_path = [self.get_mesh_path_train(train_mesh_dir, i) for i in valid_train_indices]\n",
    "            self.all_train_press_path = [self.get_pressure_data_path(train_centroid_dir, i) for i in valid_train_indices]\n",
    "            self.all_val_mesh_path = [self.get_mesh_path_train(train_mesh_dir, i) for i in valid_val_indices]\n",
    "            self.all_val_press_path = [self.get_pressure_data_path(train_centroid_dir, i) for i in valid_val_indices]\n",
    "            \n",
    "            self.all_train_centroid_path = [self.get_centroid_data_path(train_centroid_dir, i) for i in valid_train_indices]\n",
    "            self.all_val_centroid_path = [self.get_centroid_data_path(train_centroid_dir, i) for i in valid_val_indices]\n",
    "            \n",
    "            self.all_train_area_path = [self.get_area_data_path(train_centroid_dir, i) for i in valid_train_indices]\n",
    "            self.all_val_area_path = [self.get_area_data_path(train_centroid_dir, i) for i in valid_val_indices]     \n",
    "            print(f\"TrainData is {n_train-val_num}\\n ValData is {val_num}\")\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if self.mode == 'train':\n",
    "            data_dict = {\n",
    "                'vertices': self.load_centroid(self.all_train_centroid_path[idx]),\n",
    "                'pressure': self.load_pressure(self.all_train_press_path[idx]),\n",
    "                'area': self.load_area(self.all_train_area_path[idx]),\n",
    "            }\n",
    "        elif self.mode == 'val':\n",
    "            data_dict = {\n",
    "                'vertices': self.load_centroid(self.all_val_centroid_path[idx]),\n",
    "                'pressure': self.load_pressure(self.all_val_press_path[idx]),\n",
    "                'area': self.load_area(self.all_val_area_path[idx]),\n",
    "            }\n",
    "        else:\n",
    "            data_dict = {\n",
    "                'vertices': self.load_centroid(self.all_test_centroid_path[idx]),\n",
    "                'area': self.load_area(self.all_test_area_path[idx]),\n",
    "                # no presure, need to predict\n",
    "            }\n",
    "        return data_dict\n",
    "\n",
    "    def __len__(self):\n",
    "        if self.mode == 'train':\n",
    "            return len(self.all_train_mesh_path)\n",
    "        elif self.mode == 'val':\n",
    "            return len(self.all_val_mesh_path)\n",
    "        else:\n",
    "            return len(self.all_test_mesh_path)\n",
    "\n",
    "    def load_pressure(self, data_path: Path) -> np.ndarray:\n",
    "        assert data_path.exists(), \"Pressure data does not exist\"\n",
    "        press = np.load(data_path).reshape((-1,)).astype(np.float32)\n",
    "        return press\n",
    "\n",
    "    def load_centroid(self, data_path: Path) -> np.ndarray:\n",
    "        # print(data_path)\n",
    "        assert data_path.exists(), \"Pressure data does not exist\"\n",
    "        centroid = np.load(data_path).astype(np.float32)\n",
    "        return centroid\n",
    "\n",
    "    def load_area(self, data_path):\n",
    "        assert data_path.exists(), \"area data does not exist\"\n",
    "        area = np.load(data_path).reshape((-1,)).astype(np.float32)\n",
    "        return area\n",
    "\n",
    "    def get_mesh_path(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / (\"mesh_\" + str(mesh_ind).zfill(4) + \".ply\")\n",
    "\n",
    "    def get_mesh_path_train(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / 'train_mesh_0603' / (\"mesh_\" + str(mesh_ind).zfill(4) + \".ply\")\n",
    "\n",
    "    def get_mesh_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / 'Inference' / (\"mesh_\" + str(mesh_ind) + \".ply\")\n",
    "\n",
    "    def get_pressure_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / (\"press_\" + str(mesh_ind).zfill(4) + \".npy\")\n",
    "\n",
    "    def get_area_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / (\"area_\" + str(mesh_ind).zfill(4) + \".npy\")\n",
    "\n",
    "    def get_area_data_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / 'Auxiliary' / (\"area_\" + str(mesh_ind) + \".npy\")\n",
    "\n",
    "    def get_centroid_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / (\"centroid_\" + str(mesh_ind).zfill(4) + \".npy\")\n",
    "\n",
    "    def get_centroid_data_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n",
    "        return data_dir / 'Inference' / (\"centroid_\" + str(mesh_ind) + \".npy\")\n",
    "\n",
    "    def load_mesh(self, mesh_path: Path) -> o3d.t.geometry.TriangleMesh:\n",
    "        assert mesh_path.exists(), \"Mesh path does not exist\"\n",
    "        mesh = o3d.io.read_triangle_mesh(str(mesh_path))\n",
    "        mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)\n",
    "        return mesh\n",
    "\n",
    "    def vertices_from_mesh(self, mesh_path: Path):\n",
    "        mesh = self.load_mesh(mesh_path)\n",
    "        vertices = mesh.vertex.positions.numpy()\n",
    "        return vertices\n",
    "\n",
    "    def get_all_mesh_indices(\n",
    "            self, mesh_data_dir\n",
    "    ):\n",
    "        mesh_ind = []\n",
    "        all_mesh_files = os.listdir(str(mesh_data_dir) + '/train_mesh_0603')\n",
    "        all_mesh_files.sort()\n",
    "        for mesh_file_name in all_mesh_files:\n",
    "            mesh_idx = mesh_file_name.split('_')[1].split('.')[0]\n",
    "            mesh_ind.append(mesh_idx)\n",
    "        return mesh_ind\n",
    "\n",
    "    def get_valid_val_indices(self):\n",
    "        return self.valid_val_indices\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 存放损失函数的位置\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "\n",
    "# 定义损失函数\n",
    "class LpLoss(nn.Module):\n",
    "    def __init__(self, d=2, p=2, size_average=True, reduction='mean'):\n",
    "        super(LpLoss, self).__init__()\n",
    "        assert d > 0 and p > 0\n",
    "        self.d = d\n",
    "        self.p = p\n",
    "        self.reduction = reduction\n",
    "        self.size_average = size_average\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        num_examples = x.size(0)\n",
    "        h = 1.0 / (x.size(1) - 1.0)\n",
    "        all_norms = (h ** (self.d / self.p)) * torch.norm(\n",
    "            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, dim=1\n",
    "        )\n",
    "        if self.reduction == 'mean':\n",
    "            return all_norms.mean() if self.size_average else all_norms.sum()\n",
    "        elif self.reduction == 'sum':\n",
    "            return all_norms.sum()\n",
    "        else:\n",
    "            return all_norms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transolver模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "from timm.models.layers import trunc_normal_\n",
    "from einops import rearrange, repeat\n",
    "\n",
    "ACTIVATION = {'gelu': nn.GELU, 'tanh': nn.Tanh, 'sigmoid': nn.Sigmoid, 'relu': nn.ReLU, 'leaky_relu': nn.LeakyReLU(0.1),\n",
    "              'softplus': nn.Softplus, 'ELU': nn.ELU, 'silu': nn.SiLU}\n",
    "\n",
    "\n",
    "class Physics_Attention_Irregular_Mesh(nn.Module):\n",
    "    def __init__(self, dim, heads=8, dim_head=64, dropout=0., slice_num=64):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head * heads\n",
    "        self.dim_head = dim_head\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)\n",
    "\n",
    "        self.in_project_x = nn.Linear(dim, inner_dim)\n",
    "        self.in_project_fx = nn.Linear(dim, inner_dim)\n",
    "        self.in_project_slice = nn.Linear(dim_head, slice_num)\n",
    "        for l in [self.in_project_slice]:\n",
    "            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization\n",
    "        self.to_q = nn.Linear(dim_head, dim_head, bias=False)\n",
    "        self.to_k = nn.Linear(dim_head, dim_head, bias=False)\n",
    "        self.to_v = nn.Linear(dim_head, dim_head, bias=False)\n",
    "        self.to_out = nn.Sequential(\n",
    "            nn.Linear(inner_dim, dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # B N C\n",
    "        # print(x.shape)\n",
    "        _, B, N, C = x.shape\n",
    "        \n",
    "        ### (1) Slice\n",
    "        fx_mid = self.in_project_fx(x).reshape(B, N, self.heads, self.dim_head) \\\n",
    "            .permute(0, 2, 1, 3).contiguous()  # B H N C\n",
    "        x_mid = self.in_project_x(x).reshape(B, N, self.heads, self.dim_head) \\\n",
    "            .permute(0, 2, 1, 3).contiguous()  # B H N C\n",
    "        slice_weights = self.softmax(self.in_project_slice(x_mid) / self.temperature)  # B H N G\n",
    "        slice_norm = slice_weights.sum(2)  # B H G\n",
    "        slice_token = torch.einsum(\"bhnc,bhng->bhgc\", fx_mid, slice_weights)\n",
    "        slice_token = slice_token / ((slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head))\n",
    "\n",
    "        ### (2) Attention among slice tokens\n",
    "        q_slice_token = self.to_q(slice_token)\n",
    "        k_slice_token = self.to_k(slice_token)\n",
    "        v_slice_token = self.to_v(slice_token)\n",
    "        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale\n",
    "        attn = self.softmax(dots)\n",
    "        attn = self.dropout(attn)\n",
    "        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D\n",
    "\n",
    "        ### (3) Deslice\n",
    "        out_x = torch.einsum(\"bhgc,bhng->bhnc\", out_slice_token, slice_weights)\n",
    "        out_x = rearrange(out_x, 'b h n d -> b n (h d)')\n",
    "        return self.to_out(out_x)\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act='gelu', res=True):\n",
    "        super(MLP, self).__init__()\n",
    "\n",
    "        if act in ACTIVATION.keys():\n",
    "            act = ACTIVATION[act]\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "        self.n_input = n_input\n",
    "        self.n_hidden = n_hidden\n",
    "        self.n_output = n_output\n",
    "        self.n_layers = n_layers\n",
    "        self.res = res\n",
    "        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act())\n",
    "        self.linear_post = nn.Linear(n_hidden, n_output)\n",
    "        self.linears = nn.ModuleList([nn.Sequential(nn.Linear(n_hidden, n_hidden), act()) for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear_pre(x)\n",
    "        for i in range(self.n_layers):\n",
    "            if self.res:\n",
    "                x = self.linears[i](x) + x\n",
    "            else:\n",
    "                x = self.linears[i](x)\n",
    "        x = self.linear_post(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Transolver_block(nn.Module):\n",
    "    \"\"\"Transformer encoder block.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            num_heads: int,\n",
    "            hidden_dim: int,\n",
    "            dropout: float,\n",
    "            act='gelu',\n",
    "            mlp_ratio=4,\n",
    "            last_layer=False,\n",
    "            out_dim=1,\n",
    "            slice_num=32,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.last_layer = last_layer\n",
    "        self.ln_1 = nn.LayerNorm(hidden_dim)\n",
    "        self.Attn = Physics_Attention_Irregular_Mesh(hidden_dim, heads=num_heads, dim_head=hidden_dim // num_heads,\n",
    "                                                     dropout=dropout, slice_num=slice_num)\n",
    "        self.ln_2 = nn.LayerNorm(hidden_dim)\n",
    "        self.mlp = MLP(hidden_dim, hidden_dim * mlp_ratio, hidden_dim, n_layers=0, res=False, act=act)\n",
    "        if self.last_layer:\n",
    "            self.ln_3 = nn.LayerNorm(hidden_dim)\n",
    "            self.mlp2 = nn.Linear(hidden_dim, out_dim)\n",
    "\n",
    "    def forward(self, fx):\n",
    "        fx = self.Attn(self.ln_1(fx)) + fx\n",
    "        fx = self.mlp(self.ln_2(fx)) + fx\n",
    "        if self.last_layer:\n",
    "            return self.mlp2(self.ln_3(fx))\n",
    "        else:\n",
    "            return fx\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self,\n",
    "                 space_dim=1,\n",
    "                 n_layers=5,\n",
    "                 n_hidden=256,\n",
    "                 dropout=0,\n",
    "                 n_head=8,\n",
    "                 act='gelu',\n",
    "                 mlp_ratio=1,\n",
    "                 fun_dim=1,\n",
    "                 out_dim=1,\n",
    "                 slice_num=32,\n",
    "                 ref=8,\n",
    "                 unified_pos=False\n",
    "                 ):\n",
    "        super(Model, self).__init__()\n",
    "        self.__name__ = 'Transolver'\n",
    "        self.ref = ref\n",
    "        self.unified_pos = unified_pos\n",
    "        if self.unified_pos:\n",
    "            self.preprocess = MLP(fun_dim + self.ref * self.ref * self.ref, n_hidden * 2, n_hidden, n_layers=0,\n",
    "                                  res=False, act=act)\n",
    "        else:\n",
    "            self.preprocess = MLP(fun_dim + space_dim, n_hidden * 2, n_hidden, n_layers=0, res=False, act=act)\n",
    "\n",
    "        self.n_hidden = n_hidden\n",
    "        self.space_dim = space_dim\n",
    "\n",
    "        self.blocks = nn.ModuleList([Transolver_block(num_heads=n_head, hidden_dim=n_hidden,\n",
    "                                                      dropout=dropout,\n",
    "                                                      act=act,\n",
    "                                                      mlp_ratio=mlp_ratio,\n",
    "                                                      out_dim=out_dim,\n",
    "                                                      slice_num=slice_num,\n",
    "                                                      last_layer=(_ == n_layers - 1))\n",
    "                                     for _ in range(n_layers)])\n",
    "        self.initialize_weights()\n",
    "        self.placeholder = nn.Parameter((1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float))\n",
    "\n",
    "    def initialize_weights(self):\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            trunc_normal_(m.weight, std=0.02)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "\n",
    "    def get_grid(self, my_pos):\n",
    "        # my_pos 1 N 3\n",
    "        batchsize = my_pos.shape[0]\n",
    "\n",
    "        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)\n",
    "        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat([batchsize, 1, self.ref, self.ref, 1])\n",
    "        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)\n",
    "        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat([batchsize, self.ref, 1, self.ref, 1])\n",
    "        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)\n",
    "        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat([batchsize, self.ref, self.ref, 1, 1])\n",
    "        grid_ref = torch.cat((gridx, gridy, gridz), dim=-1).cuda().reshape(batchsize, self.ref ** 3, 3)  # B 4 4 4 3\n",
    "\n",
    "        pos = torch.sqrt(\n",
    "            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,\n",
    "                      dim=-1)). \\\n",
    "            reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref).contiguous()\n",
    "        return pos\n",
    "\n",
    "    def forward(self, data):\n",
    "        cfd_data = data\n",
    "        x, fx = cfd_data, None\n",
    "        x = x[None, :, :]\n",
    "        # 没有unified_pos\n",
    "        if self.unified_pos:\n",
    "            new_pos = self.get_grid(cfd_data.pos[None, :, :])\n",
    "            x = torch.cat((x, new_pos), dim=-1)\n",
    "\n",
    "        if fx is not None:\n",
    "            fx = torch.cat((x, fx), -1)\n",
    "            fx = self.preprocess(fx)\n",
    "        else:\n",
    "            fx = self.preprocess(x)\n",
    "            fx = fx + self.placeholder[None, None, :]\n",
    "\n",
    "        for block in self.blocks:\n",
    "            fx = block(fx)\n",
    "\n",
    "        return fx[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载超参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\"\"\"\n",
    "trackA 和 trackB 的超参数不同 学习率和最大学习率是一样的\n",
    "trackA: --space_dim = 3, --n_layers = 7, --n_hidden = 256, --mlp_ratio = 2, --epochs = 300\n",
    "trackB: --space_dim = 4, --n_layers = 5, --n_hidden = 256, --mlp_ratio = 1, --epochs = 155\n",
    "\"\"\"\n",
    "def get_args(epochs=300, log='train', batch=1, n_hidden=256, n_layers=7, mlp_ratio=2, space_dim=3, output_file='B_Results'):\n",
    "    parser = argparse.ArgumentParser(description=\"use config param for Transovler 7.11\")\n",
    "\n",
    "    # --------------------初始化各个参数------------------------------\n",
    "    # training Setting\n",
    "    parser.add_argument('--epochs', type=int, default=epochs, help=\"set the train epochs\")\n",
    "    parser.add_argument(\"--lr\", type=float, default=0.001, help=\"maybe 0.001 is large\")\n",
    "    parser.add_argument(\"--seed\", type=int, default=1225, help=\"Birthday\")\n",
    "    parser.add_argument(\"--log_dir\", default=log, help=\"setting the save file name\")\n",
    "    parser.add_argument(\"--batch\", type=int, default=batch, help=\"for training\")\n",
    "    parser.add_argument(\"--train_ratio\", type=float, default=0.98, help=\"500*0.98 -> 490:10 for trackA\")\n",
    "    parser.add_argument(\"--val_ratio\", type=float,default=0.2, help='for trackB')\n",
    "    parser.add_argument(\"--val_batch\", type=int, default=4, help=\"default=1\")\n",
    "    parser.add_argument(\"--max_lr\", type=float, default=0.001, help=\"use OneCycle method\")\n",
    "    # model Setting\n",
    "    parser.add_argument(\"--n_layers\", type=int, default=n_layers, help=\"maybe 5 is best for track B\")\n",
    "    parser.add_argument(\"--n_hidden\", type=int, default=n_hidden, help=\"\")\n",
    "    parser.add_argument(\"--space_dim\", type=int, default=space_dim, help=\"input data's dim, in trackB is 4\")\n",
    "    parser.add_argument(\"--mlp_ratio\", type=int, default=mlp_ratio, help='1 is best? maybe for track B')\n",
    "    parser.add_argument(\"--fun_dim\", type=int, default=0, help=\"set 0 !!\")\n",
    "    parser.add_argument(\"--n_head\", type=int, default=8, help='atten head')\n",
    "    parser.add_argument(\"--slice_num\", type=int, default=32, help='physic atten slice')\n",
    "    parser.add_argument(\"--frezze\", action='store_true', help='frezze some params')\n",
    "    parser.add_argument('--fre_layers', type=list, default=None, help='choose fre some layer')\n",
    "    parser.add_argument('--fre_epoch', type=int, default=300, help='star frezee')\n",
    "    parser.add_argument(\"--act\", default='gelu', help=\"activate function\", choices=['gelu', 'tanh', 'sigmoid', 'relu', 'leaky_relu', 'softplus', 'ELU', 'silu'])\n",
    "\n",
    "    # Test Setting\n",
    "    parser.add_argument('--result', default='best', choices=['best', 'last'], help=\"use best or last weights\")\n",
    "    parser.add_argument('--output_file', default=\"B_Results\", help='need to submit')\n",
    "    parser.add_argument('--is_zip', default=1, choices=[0, 1], help='zip or not zip')\n",
    "    parser.add_argument('--test_path', default=r'Dataset/Testset_track_A/Inference/', help='TestSet Path')\n",
    "\n",
    "    opt = parser.parse_known_args()[0]\n",
    "    return opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试 A and B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------测试结果-------------------------------\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "def testA(opt, flag='test'):\n",
    "    log = opt.log_dir\n",
    "    mode = opt.result\n",
    "    output_file = opt.output_file\n",
    "    # 检查gpu是否可用\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda\")\n",
    "        print(\"GPU is available\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "        print(\"GPU is not available, using CPU\")\n",
    "\n",
    "    # 加载模型，进行测试\n",
    "    n_layers = opt.n_layers\n",
    "    n_hidden = opt.n_hidden\n",
    "    space_dim = opt.space_dim\n",
    "    mlp_ratio = opt.mlp_ratio\n",
    "    fun_dim = opt.fun_dim\n",
    "    act = opt.act\n",
    "    slice_num = opt.slice_num\n",
    "    n_head = opt.n_head\n",
    "    model = Model(\n",
    "        n_layers=n_layers, \n",
    "        n_hidden=n_hidden, \n",
    "        space_dim=space_dim, \n",
    "        mlp_ratio=mlp_ratio, \n",
    "        fun_dim=fun_dim,\n",
    "        act=act,\n",
    "        n_head=n_head,\n",
    "        slice_num=slice_num,\n",
    "        )\n",
    "    \n",
    "    \n",
    "    # os.system(f\"mkdir -p runs_trackA/{log}/{output_file}\")\n",
    "    # 直接在当前路径下生成一个\n",
    "    os.system(f'mkdir -p {output_file}')\n",
    "\n",
    "    # 加载最优的权重模型\n",
    "    if mode == 'best':\n",
    "        if flag == 'train':\n",
    "            best_weights_path = f'runs_trackA/{log}/best_weights.pt'\n",
    "        else:\n",
    "            best_weights_path = r'best_weightsA.pt'  # 直接在当前目录下读取最佳权重\n",
    "        model.load_state_dict(torch.load(best_weights_path))\n",
    "    \n",
    "      \n",
    "    # 加载到GPU上\n",
    "    model.to(device)\n",
    "\n",
    "    \n",
    "    model.eval()\n",
    "    # 获取测试集的路径\n",
    "    test_path = read_test(path=opt.test_path)\n",
    "    # 加载测试集数据\n",
    "    dataset = TestDataset(test_dir=test_path)\n",
    "    test_loader = DataLoader(dataset, batch_size=1, shuffle=False)\n",
    "\n",
    "    # 进行测试\n",
    "    for X_batch, p in test_loader:\n",
    "\n",
    "        # 将数据部署到GPU上\n",
    "        X_batch= X_batch.to(device)        \n",
    "        outputs = model(X_batch)\n",
    "\n",
    "        # 将输出到cpu上\n",
    "        outputs = outputs.to('cpu')\n",
    "\n",
    "        # print(type(outputs))\n",
    "        # 数据进行展平，并转换为numpy的格式\n",
    "        outputs = outputs.detach().numpy().flatten()\n",
    "        print(np.max(outputs))\n",
    "\n",
    "        # 处理保存数据的路径 \n",
    "        number = str(p).split(\".\")[0][-3:]\n",
    "        path = f'{output_file}/press_{number}.npy'\n",
    "\n",
    "        # 保存数据\n",
    "        np.save(path,outputs)\n",
    "\n",
    "    print(\"-----完成数据的保存啦～=～-------\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "from torch.utils.data import  DataLoader\n",
    "from pathlib import Path\n",
    "import re\n",
    "\n",
    "# 提取数字\n",
    "def extract_numbers(s):\n",
    "    return [int(digit) for digit in re.findall(r'\\d+', s)]\n",
    "\n",
    "\n",
    "def testB(opt, flag='test'):\n",
    "    \"\"\"\n",
    "    track B's best model parser is space_dim=4 mlp_ratio=1 n_layer=5 hidden=256\n",
    "    \"\"\"\n",
    "    log = opt.log_dir\n",
    "    mode = opt.result\n",
    "    output_file = opt.output_file\n",
    "    # 检查gpu是否可用\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda\")\n",
    "        print(\"GPU is available\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "        print(\"GPU is not available, using CPU\")\n",
    "\n",
    "    # 加载模型，进行测试\n",
    "    n_layers = opt.n_layers\n",
    "    n_hidden = opt.n_hidden\n",
    "    space_dim = opt.space_dim\n",
    "    mlp_ratio = opt.mlp_ratio\n",
    "    fun_dim = opt.fun_dim\n",
    "    act = opt.act\n",
    "    slice_num = opt.slice_num\n",
    "    n_head = opt.n_head\n",
    "    model = Model(\n",
    "        n_layers=n_layers, \n",
    "        n_hidden=n_hidden, \n",
    "        space_dim=space_dim, \n",
    "        mlp_ratio=mlp_ratio, \n",
    "        fun_dim=fun_dim,\n",
    "        act=act,\n",
    "        n_head=n_head,\n",
    "        slice_num=slice_num,\n",
    "        )\n",
    "    \n",
    "    \n",
    "    os.system(f\"mkdir -p {output_file}\")\n",
    "\n",
    "    # 加载最优的权重模型\n",
    "    if mode == 'best':\n",
    "        if flag == 'train':\n",
    "            best_weights_path = f'runs_trackB/{log}/best_weights.pt'\n",
    "            \n",
    "        else:\n",
    "            best_weights_path = r'best_weightsB.pt'  # 直接在当前目录下读取最佳权重\n",
    "        model.load_state_dict(torch.load(best_weights_path))\n",
    "    \n",
    "      \n",
    "    # 加载到GPU上\n",
    "    model.to(device)\n",
    "\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    # 使用相对路径即可\n",
    "    train_dir = Path(r'./mesh_B_0603')\n",
    "    test_dir = Path(r'./Dataset/Testset_track_B/')\n",
    "    train_cent_dir = Path(r'./train_track_B')\n",
    "    test_aux_dir = Path(r'./Dataset/Testset_track_B/Auxiliary')\n",
    "\n",
    "    # 测试数据集\n",
    "    CARdataset_test = MeshPressDatasetB(train_dir, train_cent_dir, test_mesh_dir=test_dir, mode='test')\n",
    "\n",
    "    test_loader = DataLoader(CARdataset_test,\n",
    "                             batch_size=1, shuffle=False,num_workers=4,\n",
    "                             )\n",
    "    \n",
    "    \n",
    "    # 进行测试\n",
    "    for i, data in enumerate(test_loader):\n",
    "        query_points, area = data['vertices'], data['area'].unsqueeze(-1)\n",
    "        X_batch = torch.cat([query_points, area], dim=-1)\n",
    "        del query_points, area, data\n",
    "\n",
    "        # 将数据部署到GPU上\n",
    "        X_batch= X_batch.to(device)\n",
    "        \n",
    "        outputs = model(X_batch)\n",
    "\n",
    "        # 将输出到cpu上\n",
    "        outputs = outputs.squeeze(-1)\n",
    "        outputs = outputs.unsqueeze(-1).squeeze(0)\n",
    "        # outputs = outputs.cpu().numpy().flatten()\n",
    "        \n",
    "        outputs = outputs.to('cpu')\n",
    "        outputs = outputs.detach().numpy().flatten()\n",
    "        p = CARdataset_test.all_test_mesh_path[i]\n",
    "        \n",
    "        index = extract_numbers(p.name)[0]\n",
    "        index = str(index)\n",
    "\n",
    "        # 保存数据\n",
    "        np.save(f'{output_file}/press_{index}.npy',outputs)\n",
    "\n",
    "    print(\"-----完成数据的保存啦～=～-------\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练A和B （如果有必要）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练A train A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------训练模型------------------------------\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch\n",
    "from torch.utils.data import random_split,DataLoader\n",
    "from time import time\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import random\n",
    "import csv\n",
    "\n",
    "\n",
    "\n",
    "def set_seed(seed=1225):\n",
    "    # 固定随机种子\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    msg =f\"\"\"\n",
    "    =============================================\n",
    "        seed = {seed}\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "    =============================================\n",
    "    \"\"\"\n",
    "    print(msg)\n",
    "    \n",
    "def judge_file(folder_path='runs_trackA'):\n",
    "    folder_path = folder_path\n",
    "\n",
    "    if os.path.exists(folder_path) and os.path.isdir(folder_path):\n",
    "        print(\"文件夹存在\")\n",
    "    else:\n",
    "        os.system(f\"mkdir -p {folder_path}\")\n",
    "\n",
    "def trainA(opt):\n",
    "\n",
    "    log_p = opt.log_dir\n",
    "    # 设置随机种子\n",
    "    set_seed(opt.seed)\n",
    "    train_split = opt.train_ratio\n",
    "    train_batch = opt.batch\n",
    "    val_batch = opt.val_batch\n",
    "    num_epoch = opt.epochs\n",
    "    n_layers = opt.n_layers\n",
    "    n_hidden = opt.n_hidden\n",
    "    space_dim = opt.space_dim\n",
    "    mlp_ratio = opt.mlp_ratio\n",
    "    fun_dim = opt.fun_dim\n",
    "    lr = opt.lr\n",
    "    max_lr = opt.max_lr\n",
    "    act = opt.act\n",
    "    slice_num = opt.slice_num\n",
    "    n_head = opt.n_head\n",
    "    \n",
    "    # 创建保存权重的文件夹\n",
    "    save_file = 'runs_trackA'\n",
    "    judge_file(save_file)\n",
    "    log_file = os.path.join(save_file, log_p)\n",
    "    judge_file(log_file)\n",
    "\n",
    "    # 保存训练过程到csv文件中\n",
    "    f = open(f'{log_file}/train_process.csv', mode='w')\n",
    "    fc = csv.writer(f)\n",
    "    fc.writerow([\"train_loss\", \"val_loss\", 'lr'])\n",
    "\n",
    "    with open(f'{log_file}/args.txt', mode='w') as fargs:\n",
    "        fargs.write(str(opt))\n",
    "    \n",
    "    print(f\"saving args to {log_file}\")\n",
    "    \n",
    "    # 1. -加载数据集 百度给的数据集\n",
    "    feature_paths, label_paths = load_train()\n",
    "    dataset = MeshPressDataset(feature_paths, label_paths)\n",
    "    \n",
    "\n",
    "    # 检查gpu是否可用\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda\")\n",
    "        print(\"GPU is available\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "        print(\"GPU is not available, using CPU\")\n",
    "\n",
    "    \n",
    "    # 2. -划分数据集为训练集和验证集 0.8 : 0.2\n",
    "    # 可以试一下 0.9:0.1 可能效果会好很多\n",
    "    train_size = int(train_split * len(dataset))\n",
    "    val_size = len(dataset) - train_size\n",
    "    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n",
    "\n",
    "\n",
    "\n",
    "    # 3. -创建数据集加载器\n",
    "    train_loader = DataLoader(train_dataset, batch_size=train_batch, shuffle=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=val_batch, shuffle=False)\n",
    "\n",
    "\n",
    "    # 4. -实例化模型\n",
    "    model = Model(\n",
    "        n_layers=n_layers, \n",
    "        n_hidden=n_hidden, \n",
    "        space_dim=space_dim, \n",
    "        mlp_ratio=mlp_ratio, \n",
    "        fun_dim=fun_dim,\n",
    "        act=act,\n",
    "        n_head=n_head,\n",
    "        slice_num=slice_num\n",
    "        )\n",
    "\n",
    "        \n",
    "    \n",
    "    model.to(device)\n",
    "\n",
    "    params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    print(f\"---模型的参数为：{params} ---\")\n",
    "    \n",
    "\n",
    "    # 5. -定义损失函数\n",
    "    criterion = LpLoss(size_average=True)\n",
    "\n",
    "    # 6. -定义优化器 修改不同的参数，就不能加载预训练权重，会重新开始训练，最好的方法就是修改存储weights的文件夹的名称 学习率不能太大，太难了乱跑 学习率需要小一些 下降的速度也需要慢一点，太快了容易提前收敛了\n",
    "    #  之后需要修改lr lr_decay,获取弄的比较好的效果\n",
    "    # optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
    "    \n",
    "    # 7. -开始训练\n",
    "    best_val_loss = 1000\n",
    "    best_train_loss = 1000\n",
    "    # 把学习率下降调小\n",
    "\n",
    "    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "        optimizer,\n",
    "        max_lr=max_lr,\n",
    "        epochs=num_epoch,\n",
    "        div_factor=1e4, pct_start=0.25, final_div_factor=1e4,\n",
    "        steps_per_epoch=len(train_loader)\n",
    "    )\n",
    "    \n",
    "    # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.000001, max_lr=0.0008, step_size_up=800, mode='triangular')\n",
    "    grad_clip = 1000.0\n",
    "    for epoch in range(num_epoch):\n",
    "        t1 = time()\n",
    "        \n",
    "        \n",
    "        model.train()\n",
    "        train_loss = 0\n",
    "        for X_batch, y_batch in tqdm(train_loader):\n",
    "            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            outputs = model(X_batch)\n",
    "            loss = criterion(outputs,y_batch)\n",
    "            \n",
    "            loss.backward()\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            \n",
    "            \n",
    "            train_loss += loss.item()\n",
    "            # break\n",
    "        train_loss /= len(train_loader)\n",
    "\n",
    "        model.eval()\n",
    "        val_loss = 0\n",
    "        with torch.no_grad():\n",
    "            for X_batch, y_batch in val_loader:\n",
    "                X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "                outputs = model(X_batch)\n",
    "                loss = criterion(outputs, y_batch)\n",
    "                val_loss += loss.item()\n",
    "        val_loss /= len(val_loader)\n",
    "\n",
    "        \n",
    "\n",
    "        # 8. -保存模型最后一轮权重\n",
    "        last_weights_path = f'{log_file}/last_weights.pt'\n",
    "        torch.save(model.state_dict(),last_weights_path)\n",
    "\n",
    "        # 9. -保存最优权重\n",
    "        if val_loss <= best_val_loss and train_loss <= best_train_loss:\n",
    "            best_val_loss = val_loss\n",
    "            best_train_loss = train_loss\n",
    "            best_weight_path = f'{log_file}/best_weights.pt'\n",
    "            torch.save(model.state_dict(), best_weight_path)\n",
    "        \n",
    "        current_lr = optimizer.param_groups[0]['lr']\n",
    "\n",
    "        \n",
    "        if epoch+1 == opt.fre_epoch:\n",
    "            if opt.frezze:\n",
    "                # if opt.fre_epoch <= opt.epochs:\n",
    "                \n",
    "                for param in model.parameters():\n",
    "                    param.requires_grad = True\n",
    "                \n",
    "                for fre_layer in opt.fre_layers:\n",
    "                    # print(fre_layer)\n",
    "                    # 冻结选择的层数\n",
    "                    for param in model.blocks[int(fre_layer)].parameters():\n",
    "                        param.requires_grad = False\n",
    "                optimizer = torch.optim.AdamW(model.parameters(), lr=current_lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
    "                scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "                    optimizer,\n",
    "                    max_lr=current_lr,\n",
    "                    epochs=num_epoch-opt.fre_epoch,\n",
    "                    div_factor=1e4, pct_start=0.001, final_div_factor=1e4,\n",
    "                    steps_per_epoch=len(train_loader)\n",
    "                )\n",
    "        \n",
    "        t2 =time()\n",
    "        t = t2 - t1\n",
    "        print(f'Epoch {epoch+1}/{num_epoch}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Current_lr: {current_lr}, Spend Time: {round(t,2)}s')\n",
    "        fc.writerow([train_loss, val_loss, current_lr])\n",
    "\n",
    "    f.close()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练B train B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练模型\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from pathlib import Path\n",
    "from time import time\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import random\n",
    "import csv\n",
    "\n",
    "\n",
    "def set_seed(seed=1225):\n",
    "    \n",
    "    # 固定随机种子\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    msg =f\"\"\"\n",
    "    =============================================\n",
    "        seed = {seed}\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "    =============================================\n",
    "    \"\"\"\n",
    "    print(msg)\n",
    "    \n",
    "def judge_file(folder_path='runs_trackA'):\n",
    "    folder_path = folder_path\n",
    "\n",
    "    if os.path.exists(folder_path) and os.path.isdir(folder_path):\n",
    "        print(\"文件夹存在\")\n",
    "    else:\n",
    "        os.system(f\"mkdir -p {folder_path}\")\n",
    "\n",
    "def trainB(opt):\n",
    "\n",
    "    log_p = opt.log_dir\n",
    "    # 设置随机种子\n",
    "    set_seed(opt.seed)\n",
    "    train_split = opt.train_ratio\n",
    "    train_batch = opt.batch\n",
    "    val_batch = opt.val_batch\n",
    "    num_epoch = opt.epochs\n",
    "    n_layers = opt.n_layers\n",
    "    n_hidden = opt.n_hidden\n",
    "    space_dim = opt.space_dim\n",
    "    mlp_ratio = opt.mlp_ratio\n",
    "    fun_dim = opt.fun_dim\n",
    "    lr = opt.lr\n",
    "    max_lr = opt.max_lr\n",
    "    act = opt.act\n",
    "    slice_num = opt.slice_num\n",
    "    n_head = opt.n_head\n",
    "    \n",
    "    # 创建保存权重的文件夹\n",
    "    save_file = 'runs_trackB'\n",
    "    judge_file(save_file)\n",
    "    log_file = os.path.join(save_file, log_p)\n",
    "    judge_file(log_file)\n",
    "\n",
    "    # 保存训练过程到csv文件中\n",
    "    f = open(f'{log_file}/train_process.csv', mode='w')\n",
    "    fc = csv.writer(f)\n",
    "    fc.writerow([\"train_loss\", \"val_loss\", 'lr'])\n",
    "\n",
    "    with open(f'{log_file}/args.txt', mode='w') as fargs:\n",
    "        fargs.write(str(opt))\n",
    "    \n",
    "    print(f\"saving args to {log_file}\")\n",
    "    \n",
    "    \n",
    "    # 检查gpu是否可用\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda\")\n",
    "        print(\"GPU is available\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "        print(\"GPU is not available, using CPU\")\n",
    "\n",
    "    \n",
    "    train_dir = Path(r'./mesh_B_0603')\n",
    "    test_dir = Path(r'./Testset_track_B/')\n",
    "    train_cent_dir = Path(r'./train_track_B')\n",
    "    test_aux_dir = Path(r'.Testset_track_B/Auxiliary')\n",
    "\n",
    "    CARdataset_train = MeshPressDatasetB(train_dir, train_cent_dir, test_dir, val_ratio=0.2)\n",
    "    CARdataset_val = MeshPressDatasetB(train_dir, train_cent_dir, test_dir, val_ratio=0.2, mode='val')\n",
    "\n",
    "    # 默认400训练集 100测试集\n",
    "    train_loader = DataLoader(CARdataset_train,\n",
    "                              batch_size=opt.batch, shuffle=True, pin_memory=True,\n",
    "                              )\n",
    "    val_loader = DataLoader(CARdataset_val,\n",
    "                            batch_size=1, shuffle=False\n",
    "                            )\n",
    "\n",
    "    # 4. -实例化模型\n",
    "    model = Model(\n",
    "        n_layers=n_layers, \n",
    "        n_hidden=n_hidden, \n",
    "        space_dim=space_dim, \n",
    "        mlp_ratio=mlp_ratio, \n",
    "        fun_dim=fun_dim,\n",
    "        act=act,\n",
    "        n_head=n_head,\n",
    "        slice_num=slice_num\n",
    "        )\n",
    "\n",
    "    \n",
    "    \n",
    "    model.to(device)\n",
    "\n",
    "    \n",
    "    params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    print(f\"---模型的参数为：{params} ---\")\n",
    "    \n",
    "    del params\n",
    "    \n",
    "    # 5. -定义损失函数\n",
    "    criterion = LpLoss(size_average=True)\n",
    "\n",
    "    # 6. -定义优化器 修改不同的参数，就不能加载预训练权重，会重新开始训练，最好的方法就是修改存储weights的文件夹的名称 学习率不能太大，太难了乱跑 学习率需要小一些 下降的速度也需要慢一点，太快了容易提前收敛了\n",
    "    #  之后需要修改lr lr_decay,获取弄的比较好的效果\n",
    "    # optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
    "    \n",
    "    # 7. -开始训练\n",
    "    best_val_loss = 1000\n",
    "    best_train_loss = 1000\n",
    "    # 把学习率下降调小\n",
    "\n",
    "    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "        optimizer,\n",
    "        max_lr=max_lr,\n",
    "        epochs=num_epoch,\n",
    "        div_factor=1e4, pct_start=0.2, final_div_factor=1e4,\n",
    "        steps_per_epoch=len(train_loader)\n",
    "    )\n",
    "    \n",
    "    # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.000001, max_lr=0.0008, step_size_up=800, mode='triangular')\n",
    "    # del model\n",
    "    grad_clip = 1000.0\n",
    "    for epoch in range(num_epoch):\n",
    "        t1 = time()\n",
    "        \n",
    "        \n",
    "        model.train()\n",
    "        torch.cuda.empty_cache()\n",
    "        train_loss = 0\n",
    "        for data in tqdm(train_loader):\n",
    "            optimizer.zero_grad()\n",
    "                \n",
    "            query_points, y_batch, area = data['vertices'], data['pressure'], data['area'].unsqueeze(-1)\n",
    "            X_batch = torch.cat([query_points, area], dim=-1)\n",
    "            del query_points, area, data\n",
    "            # query_points, y = query_points.to(device), y.to(device),\n",
    "            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "\n",
    "            \n",
    "\n",
    "            outputs = model(X_batch)\n",
    "            loss = criterion(outputs,y_batch)\n",
    "            # del X_batch, y_batch\n",
    "            loss.backward()\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            \n",
    "            \n",
    "            train_loss += loss.item()\n",
    "            # break\n",
    "        train_loss /= len(train_loader)\n",
    "\n",
    "        model.eval()\n",
    "        val_loss = 0\n",
    "        with torch.no_grad():\n",
    "            for data in val_loader:\n",
    "                \n",
    "                query_points, y_batch, area = data['vertices'], data['pressure'], data['area'].unsqueeze(-1)\n",
    "                X_batch = torch.cat([query_points, area], dim=-1)\n",
    "                del query_points, area, data\n",
    "                \n",
    "                X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "                \n",
    "                \n",
    "                outputs = model(X_batch)\n",
    "                loss = criterion(outputs, y_batch)\n",
    "                \n",
    "                # del X_batch, y_batch\n",
    "                val_loss += loss.item()\n",
    "        val_loss /= len(val_loader)\n",
    "\n",
    "        \n",
    "\n",
    "        # 8. -保存模型最后一轮权重\n",
    "        last_weights_path = f'{log_file}/last_weights.pt'\n",
    "        torch.save(model.state_dict(),last_weights_path)\n",
    "\n",
    "        # 9. -保存最优权重\n",
    "        if val_loss <= best_val_loss and train_loss <= best_train_loss:\n",
    "            best_val_loss = val_loss\n",
    "            best_train_loss = train_loss\n",
    "            best_weight_path = f'{log_file}/best_weights.pt'\n",
    "            torch.save(model.state_dict(), best_weight_path)\n",
    "        \n",
    "        current_lr = optimizer.param_groups[0]['lr']\n",
    "        \n",
    "        if epoch+1 == opt.fre_epoch:\n",
    "            if opt.frezze:\n",
    "                # if opt.fre_epoch <= opt.epochs:\n",
    "                \n",
    "                for param in model.parameters():\n",
    "                    param.requires_grad = True\n",
    "                \n",
    "                for fre_layer in opt.fre_layers:\n",
    "                    # print(fre_layer)\n",
    "                    # 冻结选择的层数\n",
    "                    for param in model.blocks[int(fre_layer)].parameters():\n",
    "                        param.requires_grad = False\n",
    "                optimizer = torch.optim.AdamW(model.parameters(), lr=current_lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
    "                scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "                    optimizer,\n",
    "                    max_lr=current_lr,\n",
    "                    epochs=num_epoch-opt.fre_epoch,\n",
    "                    div_factor=1e4, pct_start=0.001, final_div_factor=1e4,\n",
    "                    steps_per_epoch=len(train_loader)\n",
    "                )\n",
    "        \n",
    "        t2 =time()\n",
    "        t = t2 - t1\n",
    "        print(f'Epoch {epoch+1}/{num_epoch}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Current_lr: {current_lr}, Spend Time: {round(t,2)}s')\n",
    "        fc.writerow([train_loss, val_loss, current_lr])\n",
    "\n",
    "    f.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 进行测试，并输出C结果"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "C的结果是分别使用A和B的最优权重实现的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# testA\n",
    "optA = get_args(n_hidden=256, n_layers=7, space_dim=3, mlp_ratio=2)\n",
    "testA(opt=optA)\n",
    "# testB\n",
    "optB = get_args(n_hidden=256, n_layers=5, mlp_ratio=1, space_dim=4)\n",
    "testB(opt=optB)\n",
    "import zipfile\n",
    "with zipfile.ZipFile('B_Results.zip','w') as zipObj:\n",
    "    file_list = os.listdir('B_Results')\n",
    "    for f in file_list:\n",
    "        p = os.path.join('B_Results/', f)\n",
    "        zipObj.write(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练C模型，并测试C模型 解除注释即可执行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optA = get_args(n_hidden=256, n_layers=7, space_dim=3, mlp_ratio=2, epochs=300, log='trainA', batch=2)\n",
    "# trainA(opt=optA)\n",
    "# testA(opt=optA, flag='train')\n",
    "\n",
    "# optB = get_args(n_hidden=256, n_layers=5, mlp_ratio=1, space_dim=4, epochs=155, log='trainB',batch=1)\n",
    "# trainB(opt=optB)\n",
    "# testB(opt=optB,flag='train')\n",
    "# # 保存为压缩文件\n",
    "# import zipfile\n",
    "# with zipfile.ZipFile('B_Results.zip','w') as zipObj:\n",
    "#     file_list = os.listdir('B_Results')\n",
    "#     for f in file_list:\n",
    "#         p = os.path.join('B_Results/', f)\n",
    "#         zipObj.write(p)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
